#include "Database.h"

using namespace std;

//! Function to read the scene information from the abstract xml file.
/*!
 * The xml file that is read contains the information about the database in
 * an abstract form. The function extracts the information and creates scene
 * information for every scene in the benchmark.
 * @param scenes A reference to a vector. The function will fill this vector
 * with the scene information. 
 * @param xmlFile The xml file name.
 */
void Database::readScenesInfo(vector<SceneInfo*> &scenes, string xmlFile) {
  using boost::property_tree::ptree;
  ptree pt;


  // Load the XML file
  read_xml(xmlFile, pt, boost::property_tree::xml_parser::trim_whitespace);
  
  string visgrabInPath = pt.get<string>("VisGraBScenes.VisGraBInPath");
  string visgrabSimPath = pt.get<string>("VisGraBScenes.VisGraBSimPath");

  vector< PoseCondition > poseConditions;

  int sceneID = 0;

  // Get Pose conditions
  BOOST_FOREACH(ptree::value_type &v, pt.get_child("VisGraBScenes")) {
    if((string)v.first.data() == "PoseConditions") {
      ptree ptPoseConds = (ptree) v.second;

      BOOST_FOREACH(ptree::value_type &v2, ptPoseConds) {
	if((string)v2.first.data() == "Complexity") {
	  ptree ptComplexity = (ptree) v2.second;
	  int nrObjects = ptComplexity.get<int>("<xmlattr>.nrObjects");

	  BOOST_FOREACH(ptree::value_type &v3, ptComplexity) {
	    if((string)v3.first.data() == "PoseCondition") {
	      ptree ptPoseCond = (ptree) v3.second;
	      PoseCondition poseCond;
	      poseCond.nrObjects = nrObjects;
	      poseCond.type = ptPoseCond.get<string>("<xmlattr>.type");
	      poseCond.startID = ptPoseCond.get<int>("<xmlattr>.startID");
	      poseCond.endID = ptPoseCond.get<int>("<xmlattr>.endID");
	      poseConditions.push_back(poseCond);	      
	    }
	  }
	}
      }
    }

    // Go through all texture conditions
    if((string)v.first.data() == "TextureCondition") {
      ptree ptTexCond = (ptree) v.second;
      string textureCondition = ptTexCond.get<string>("<xmlattr>.type");
      
      BOOST_FOREACH(ptree::value_type &v2, ptTexCond) {
	if((string)v2.first.data() == "Complexity") {
	  ptree ptComplexity = (ptree) v2.second;
	  int nrObjects = ptComplexity.get<int>("<xmlattr>.nrObjects");
	  
	  BOOST_FOREACH(ptree::value_type &v3, ptComplexity) {
	    if((string)v3.first.data() == "SceneSet") {
	      ptree ptSceneSet = (ptree) v3.second;
	      int sceneSetID = ptSceneSet.get<int>("<xmlattr>.id");
	      string sceneSet = ptSceneSet.get<string>("<xmlattr>.name");
	      vector<string> objects;
	      for(int objI=0; objI<nrObjects; objI++) {
		char fName[20];
		sprintf(fName,"<xmlattr>.object%d", objI+1);
		string object = ptSceneSet.get<string>(fName);
		objects.push_back(object);
	      }
	      int startID = ptSceneSet.get<int>("<xmlattr>.startID");
	      int endID = ptSceneSet.get<int>("<xmlattr>.endID");
	      
	      for(int poseI=startID; poseI<endID+1; poseI++) {
		SceneInfo *scene = new SceneInfo;
		scene->sceneID = sceneID;
		sceneID++;
		char poseIstr[10];
		sprintf(poseIstr, "img_%02d", poseI);
		string sceneLocator = textureCondition + "/" + sceneSet + "/" + poseIstr;
		scene->sceneLocator = sceneLocator;
		//scene->scenePath = visgrabInPath + "/" + sceneLocator;
		scene->textureCondition = textureCondition;
		scene->poseCondition = getPoseCondition(poseConditions, nrObjects, poseI);
		scene->nrObjects = nrObjects;
		scene->objects = objects;
		scene->sceneSetID = sceneSetID;
		scene->sceneSet = sceneSet;
		scene->poseID = poseI-1; //! NB. In the visgrab database, counting starts at 1, poseID therefore -1
		scene->leftImg = visgrabInPath + "/" + sceneLocator + "_left_rect.ppm";
		scene->rightImg = visgrabInPath + "/" + sceneLocator + "_right_rect.ppm";
		scene->disp = visgrabInPath + "/" + sceneLocator + "_disp.dat";
		scene->pointCloud = visgrabInPath + "/" + sceneLocator + "_point_cloud.dat";
		scene->wc = visgrabSimPath + "/" + textureCondition + "/workcells/" + 
		  sceneSet + "/" + poseIstr + ".wc.xml";
		scene->dwc = visgrabSimPath + "/" + textureCondition + "/workcells/" + 
		  sceneSet + "/" + poseIstr + ".dwc.xml";
		scenes.push_back(scene);
	      }
	    }
	  }
	}
      }
    }
  } 
}

//! Function to read the scene information from the xml file.
/*!
 * The xml file that is read contains the full information about the database. 
 * The function  creates scene information for every scene in the benchmark.
 * 
 * The full information file can be created from the abstract information file
 * by reading it using readScenesInfo(vector<SceneInfo*> &scenes, string xmlFile)
 * and writing it using writeDatabaseInfo(vector<SceneInfo*> &scenes, string xmlFile)
 * @param scenes A reference to a vector. The function will fill this vector
 * with the scene information. 
 * @param xmlFile The xml file name.
 */
void Database::readDatabaseInfo(vector<SceneInfo*> &scenes, string xmlFile) {
  using boost::property_tree::ptree;
  ptree pt;

  read_xml(xmlFile, pt, boost::property_tree::xml_parser::trim_whitespace);

  string benchmarkType = pt.get<string>("visgrab.benchmark.<xmlattr>.type");
  if(benchmarkType=="GraspUnknownObject") {
    BOOST_FOREACH(ptree::value_type &v, pt.get_child("visgrab.benchmark")) {
      if((string)v.first.data() == "scene") {
	ptree ptScene = (ptree) v.second;
	SceneInfo *scene = new SceneInfo;
	scene->sceneID = ptScene.get<int>("<xmlattr>.id");
	scene->sceneLocator = ptScene.get<string>("info.sceneLocator");
	scene->textureCondition = ptScene.get<string>("info.textureCondition");
	scene->poseCondition = ptScene.get<string>("info.poseCondition");
	scene->sceneSet = ptScene.get<string>("info.sceneSet");
	scene->sceneSetID = ptScene.get<int>("info.sceneSetID");
	scene->poseID = ptScene.get<int>("info.poseID");
	scene->leftImg = ptScene.get<string>("input.leftImg");
	scene->rightImg = ptScene.get<string>("input.rightImg");
	scene->disp = ptScene.get<string>("input.disp");
	scene->pointCloud = ptScene.get<string>("input.pointCloud");
	scene->wc = ptScene.get<string>("simulator.wc");
	scene->dwc = ptScene.get<string>("simulator.dwc");
	// Objects
	scene->nrObjects = ptScene.get<int>("info.nrObjects");
	BOOST_FOREACH(ptree::value_type &v, ptScene.get_child("info.objects")) {
	  scene->objects.push_back( v.second.data() );
	}
	scenes.push_back(scene);
      }
    }
  }
  std::cout << "Nr of scenes to process: " << scenes.size() << std::endl;
}

//! Function to write the full database information
/*
 * Writes the full database information to the xml file. All information of
 * every scene is written to file. The saved file can be reloaded with
 * readDatabaseInfo(vector<SceneInfo*> &scenes, string xmlFile)
 * @param scenes A reference to the vector containing all scene information. 
 * @param xmlFile The xml file name. 
 */
void Database::writeDatabaseInfo(vector<SceneInfo*> &scenes, string xmlFile) {
  using boost::property_tree::ptree;
  ptree pt;

  // pt.put("debug.level", m_level);
  pt.put("visgrab.benchmark.<xmlattr>.type", "GraspUnknownObject");
  for(unsigned int i=0; i<scenes.size(); i++) {
    ptree ptChild;
    ptChild.put("<xmlattr>.id", i);
    ptChild.put("info.sceneLocator", scenes[i]->sceneLocator);
    //ptChild.put("info.scenePath", scenes[i]->scenePath);
    ptChild.put("info.textureCondition", scenes[i]->textureCondition);
    ptChild.put("info.poseCondition", scenes[i]->poseCondition);
    ptChild.put("info.sceneSet", scenes[i]->sceneSet);
    ptChild.put("info.sceneSetID", scenes[i]->sceneSetID);
    ptChild.put("info.poseID", scenes[i]->poseID);
    ptChild.put("info.nrObjects", scenes[i]->nrObjects);
    for(unsigned int oI=0; oI<scenes[i]->objects.size(); oI++) 
      ptChild.add("info.objects.object", scenes[i]->objects[oI]); //, true);

    ptChild.put("input.leftImg", scenes[i]->leftImg);
    ptChild.put("input.rightImg", scenes[i]->rightImg);
    ptChild.put("input.disp", scenes[i]->disp);
    ptChild.put("input.pointCloud", scenes[i]->pointCloud);

    ptChild.put("simulator.wc", scenes[i]->wc);
    ptChild.put("simulator.dwc", scenes[i]->dwc);
    pt.add_child("visgrab.benchmark.scene", ptChild);//, true);
  }
  boost::property_tree::xml_writer_settings<char> settings('\t', 1);
  write_xml(xmlFile, pt, std::locale(), settings);
}

/*!
 * Get the scene data (left and right rectified images, disparity map, and point cloud).
 * Data is read in the class ScenData.
 * @param sceneInfo A pointer to the scene info
 */
SceneData* Database::getSceneData(SceneInfo* sceneInfo, std::string rootPath) {

  SceneData* sceneData = new SceneData(sceneInfo, rootPath);

  return sceneData;
}

/*!
 * Get the pose condition belonging to the current poseID.
 * @param poseConditions Reference to the vector containing all pose conditions.
 * @param nrObjects The number of objects in the scene
 * @param poseID The ID of the current pose
 * @return A string containing the pose condition.
 */
string Database::getPoseCondition(vector<PoseCondition> &poseConditions, int nrObjects, int poseID) {
  for(unsigned int i=0; i<poseConditions.size(); i++) {
    if(poseConditions[i].nrObjects == nrObjects &&
       poseConditions[i].startID <= poseID &&
       poseID <= poseConditions[i].endID)
      return poseConditions[i].type;
  }
  return "unknown";
}

/*!
 * Add the location of the collision-info file to each scene. The user can specify the location
 * of the collision-info files.
 * @param scenes A reference to the vector containing all scene information. 
 * @param collisionInfoPath A string with the collision-info directory
 */
void Database::addCollisionInfoDirectory(vector<SceneInfo*> &scenes, string collisionInfoPath) {
  for(unsigned int i=0; i<scenes.size(); i++) {
    scenes[i]->collision = collisionInfoPath + "/" + scenes[i]->sceneLocator + "_collisionInfo.dat.stl";
  }
}
